#include "KDTree.h"
#include <queue>
#include <limits>
#include <numeric>  // 添加对std::iota的支持

KDTree::KDTree(const Eigen::MatrixXd& points) 
    : points_(points), dim_(points.cols()) {
    
    std::vector<size_t> indices(points.rows());
    std::iota(indices.begin(), indices.end(), 0);
    root_ = buildTree(indices, 0);
}

std::shared_ptr<KDNode> KDTree::buildTree(
    const std::vector<size_t>& indices,
    int depth) {
    
    if (indices.empty()) return nullptr;
    
    int axis = depth % dim_;
    auto n = indices.size();
    
    // 选择中位数作为分割点
    std::vector<size_t> sorted_indices = indices;
    std::nth_element(sorted_indices.begin(), 
                    sorted_indices.begin() + n/2,
                    sorted_indices.end(),
                    [this, axis](size_t i1, size_t i2) {
                        return points_(i1, axis) < points_(i2, axis);
                    });
    
    size_t median_idx = sorted_indices[n/2];
    auto node = std::make_shared<KDNode>(points_.row(median_idx), median_idx, axis);
    
    // 递归构建左右子树
    if (n > 1) {
        std::vector<size_t> left_indices(sorted_indices.begin(), 
                                       sorted_indices.begin() + n/2);
        std::vector<size_t> right_indices(sorted_indices.begin() + n/2 + 1,
                                        sorted_indices.end());
        
        node->left = buildTree(left_indices, depth + 1);
        node->right = buildTree(right_indices, depth + 1);
    }
    
    return node;
}

std::pair<size_t, double> KDTree::nearest(const Eigen::VectorXd& query) const {
    size_t bestIndex = 0;
    double bestDist = std::numeric_limits<double>::max();
    nearestImpl(query, root_, bestIndex, bestDist);
    return {bestIndex, bestDist};
}

void KDTree::nearestImpl(const Eigen::VectorXd& query,
                        const std::shared_ptr<KDNode>& node,
                        size_t& bestIndex,
                        double& bestDist) const {
    if (!node) return;
    
    double dist = (node->point - query).squaredNorm();
    if (dist < bestDist) {
        bestDist = dist;
        bestIndex = node->index;
    }
    
    double diff = query(node->axis) - node->point(node->axis);
    std::shared_ptr<KDNode> near = diff <= 0 ? node->left : node->right;
    std::shared_ptr<KDNode> far = diff <= 0 ? node->right : node->left;
    
    nearestImpl(query, near, bestIndex, bestDist);
    
    if (diff * diff < bestDist) {
        nearestImpl(query, far, bestIndex, bestDist);
    }
}

void KDTree::kNearest(const Eigen::VectorXd& query, int k,
                     std::vector<size_t>& indices,
                     std::vector<double>& distances) const {
    
    std::vector<std::pair<double, size_t>> nearest;
    nearest.reserve(k);
    kNearestImpl(query, root_, nearest, k);
    
    // 排序并返回结果
    std::sort(nearest.begin(), nearest.end());
    indices.resize(nearest.size());
    distances.resize(nearest.size());
    
    for (size_t i = 0; i < nearest.size(); ++i) {
        distances[i] = nearest[i].first;
        indices[i] = nearest[i].second;
    }
}

void KDTree::kNearestImpl(const Eigen::VectorXd& query,
                         const std::shared_ptr<KDNode>& node,
                         std::vector<std::pair<double, size_t>>& nearest,
                         int k) const {
    if (!node) return;
    
    double dist = (node->point - query).squaredNorm();
    
    if (nearest.size() < k) {
        nearest.emplace_back(dist, node->index);
        std::push_heap(nearest.begin(), nearest.end());
    } else if (dist < nearest.front().first) {
        std::pop_heap(nearest.begin(), nearest.end());
        nearest.back() = {dist, node->index};
        std::push_heap(nearest.begin(), nearest.end());
    }
    
    double diff = query(node->axis) - node->point(node->axis);
    std::shared_ptr<KDNode> near = diff <= 0 ? node->left : node->right;
    std::shared_ptr<KDNode> far = diff <= 0 ? node->right : node->left;
    
    kNearestImpl(query, near, nearest, k);
    
    if (nearest.size() < k || diff * diff < nearest.front().first) {
        kNearestImpl(query, far, nearest, k);
    }
} 